import cv2
import dlib
import numpy as np
from imutils import face_utils
from PIL import Image, ImageDraw, ImageFont


# 加载面部关键点预测器
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
# 初始化人脸检测器
detector = dlib.get_frontal_face_detector()

# 初始化眨眼和打哈欠计数器
blink_counter = 0
yawn_counter = 0
total_frames = 0

# 调整眼睛纵横比阈值以提高灵敏度
EYE_AR_THRESH = 0.22
EYE_AR_CONSEC_FRAMES = 2

# 嘴巴纵横比阈值
MOUTH_AR_THRESH = 0.6
MOUTH_AR_CONSEC_FRAMES = 3

# 初始化连续闭眼和张嘴帧数
COUNTER = 0
TOTAL = 0
YAWN_COUNTER = 0
YAWN_TOTAL = 0


def calculate_ear(eye):
    """
    计算眼睛的纵横比 (EAR)
    :param eye: 眼睛的关键点坐标
    :return: 眼睛的纵横比
    """
    A = np.linalg.norm(eye[1] - eye[5])
    B = np.linalg.norm(eye[2] - eye[4])
    C = np.linalg.norm(eye[0] - eye[3])
    ear = (A + B) / (2.0 * C)
    return ear


def calculate_mar(mouth):
    """
    计算嘴巴的纵横比 (MAR)
    :param mouth: 嘴巴的关键点坐标
    :return: 嘴巴的纵横比
    """
    A = np.linalg.norm(mouth[2] - mouth[6])
    B = np.linalg.norm(mouth[3] - mouth[5])
    C = np.linalg.norm(mouth[0] - mouth[4])
    mar = (A + B) / (2.0 * C)
    return mar


def get_head_pose(shape):
    """
    计算头部姿态
    :param shape: 68 个面部关键点
    :return: 俯仰角、偏航角
    """
    image_points = np.array([
        shape[30],  # 鼻子尖端
        shape[8],  # 下巴
        shape[36],  # 左眼内角
        shape[45],  # 右眼内角
        shape[48],  # 左嘴角
        shape[54]  # 右嘴角
    ], dtype="double")

    model_points = np.array([
        (0.0, 0.0, 0.0),  # 鼻子尖端
        (0.0, -330.0, -65.0),  # 下巴
        (-225.0, 170.0, -135.0),  # 左眼内角
        (225.0, 170.0, -135.0),  # 右眼内角
        (-150.0, -150.0, -125.0),  # 左嘴角
        (150.0, -150.0, -125.0)  # 右嘴角
    ])

    size = (480, 640)  # 图像尺寸
    focal_length = size[1]
    center = (size[1] / 2, size[0] / 2)
    camera_matrix = np.array(
        [[focal_length, 0, center[0]],
         [0, focal_length, center[1]],
         [0, 0, 1]], dtype="double"
    )

    dist_coeffs = np.zeros((4, 1))
    (success, rotation_vector, translation_vector) = cv2.solvePnP(model_points, image_points, camera_matrix,
                                                                 dist_coeffs,
                                                                 flags=cv2.SOLVEPNP_ITERATIVE)

    (nose_end_point2D, jacobian) = cv2.projectPoints(np.array([(0.0, 0.0, 1000.0)]), rotation_vector,
                                                     translation_vector, camera_matrix, dist_coeffs)

    p1 = (int(image_points[0][0]), int(image_points[0][1]))
    p2 = (int(nose_end_point2D[0][0][0]), int(nose_end_point2D[0][0][1]))

    # 计算俯仰角和偏航角
    rotation_mat, _ = cv2.Rodrigues(rotation_vector)
    pose_mat = cv2.hconcat((rotation_mat, translation_vector))
    _, _, _, _, _, _, euler_angle = cv2.decomposeProjectionMatrix(pose_mat)

    pitch = euler_angle[0][0]
    yaw = euler_angle[1][0]
    return pitch, yaw


def calculate_focus_percentage(pitch, yaw, blink_count, yawn_count):
    """
    计算专注度百分比
    :param pitch: 头部俯仰角
    :param yaw: 头部偏航角
    :param blink_count: 眨眼次数
    :param yawn_count: 打哈欠次数
    :return: 专注度百分比
    """
    pitch_score = max(0, 1 - abs(pitch) / 30)
    yaw_score = max(0, 1 - abs(yaw) / 30)
    blink_score = max(0, 1 - blink_count / 10)
    yawn_score = max(0, 1 - yawn_count / 5)

    focus_percentage = (pitch_score + yaw_score + blink_score + yawn_score) / 4
    return focus_percentage


def get_focus_color(focus_percentage):
    """
    根据专注度百分比获取颜色
    :param focus_percentage: 专注度百分比
    :return: 颜色 (R, G, B)
    """
    if focus_percentage >= 0.8:
        return (0, 255, 0)  # 绿色
    elif focus_percentage >= 0.5:
        return (255, 255, 0)  # 黄色
    else:
        return (255, 0, 0)  # 红色


# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 使用 dlib 进行人脸检测
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = detector(gray)

    for face in faces:
        # 获取面部关键点
        shape = predictor(gray, face)
        shape = face_utils.shape_to_np(shape)

        # 提取眼睛和嘴巴的关键点
        left_eye = shape[42:48]
        right_eye = shape[36:42]
        mouth = shape[60:68]

        # 计算眼睛和嘴巴的纵横比
        left_ear = calculate_ear(left_eye)
        right_ear = calculate_ear(right_eye)
        ear = (left_ear + right_ear) / 2.0
        mar = calculate_mar(mouth)

        # 检测眨眼
        if ear < EYE_AR_THRESH:
            COUNTER += 1
        else:
            if COUNTER >= EYE_AR_CONSEC_FRAMES:
                TOTAL += 1
            COUNTER = 0

        # 检测打哈欠
        if mar > MOUTH_AR_THRESH:
            YAWN_COUNTER += 1
        else:
            if YAWN_COUNTER >= MOUTH_AR_CONSEC_FRAMES:
                YAWN_TOTAL += 1
            YAWN_COUNTER = 0

        # 计算头部姿态
        pitch, yaw = get_head_pose(shape)

        # 计算专注度百分比
        focus_percentage = calculate_focus_percentage(pitch, yaw, TOTAL, YAWN_TOTAL)

        # 获取专注度颜色
        focus_color = get_focus_color(focus_percentage)

        # 绘制面部关键点
        for (x, y) in shape:
            cv2.circle(frame, (x, y), 2, (0, 255, 0), -1)

        # 使用 PIL 显示中文信息
        img_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(img_pil)
        font = ImageFont.truetype("simhei.ttf", 20)
        draw.text((10, 30), f"眨眼次数: {TOTAL}", font=font, fill=(0, 255, 0))
        draw.text((10, 60), f"头部偏转角(上下): {pitch:.2f}°", font=font, fill=(0, 255, 0))
        draw.text((10, 90), f"头部偏转角(左右): {yaw:.2f}°", font=font, fill=(0, 255, 0))
        draw.text((10, 120), f"专注度: {focus_percentage * 100:.2f}%", font=font, fill=focus_color)
        frame = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

    # 显示帧
    cv2.imshow('Focus Detection', frame)

    # 按 'q' 键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()